# langchain_agent_demo_optimized.py：优化版 Agent 代码（支持多轮对话+无警告）
from langchain.tools import tool
from langchain_openai import ChatOpenAI  # 正确导入 ChatOpenAI
from langchain.agents import ZeroShotAgent, AgentExecutor
from langchain.prompts import PromptTemplate
from langchain.memory import ConversationBufferMemory  # 多轮对话记忆
from dotenv import load_dotenv
import os

# 1. 加载环境变量（OpenAI/高德/RAG等密钥）
load_dotenv()


# -------------------------- 工具定义 --------------------------
@tool(
    description="用于查询指定城市的当前天气。参数：city（必填，城市名，示例：北京、上海），返回值：实时温度、天气状况、风力。需调用高德天气API获取数据，不可编造。")
def WeatherTool(city: str) -> str:
    amap_key = "ec0954c5a50ea226ca13ae6a05da4e98"
    if not amap_key:
        return "获取天气失败：未配置高德API密钥（AMAP_API_KEY）"
    url = f"https://restapi.amap.com/v3/weather/weatherInfo?city={city}&key={amap_key}&extensions=base"
    try:
        import requests
        response = requests.get(url)
        data = response.json()
        if data["status"] == "1" and len(data["lives"]) > 0:
            live = data["lives"][0]
            return f"{city}当前天气：{live['weather']}，温度{live['temperature']}℃，风力{live['windpower']}级，湿度{live['humidity']}%"
        else:
            return f"获取{city}天气失败：{data.get('info', '未知错误')}"
    except Exception as e:
        return f"调用天气API失败：{str(e)}"


@tool(
    description="用于计算加减乘除运算。参数：expression（必填，数学表达式，示例：2+3*4、10-5，无空格），返回值：计算结果。仅处理数字和+、-、*、/。")
def CalculatorTool(expression: str) -> str:
    try:
        result = eval(expression)
        return f"计算结果：{expression} = {result}"
    except Exception as e:
        return f"计算失败：{str(e)}，请检查表达式（示例：2+3，不要有空格）"


@tool(
    description="用于查询公司内部知识库（如年假政策、产品价格）。参数：question（必填，查询问题，示例：公司年假政策、产品A价格），返回值：知识库答案。公司相关问题必须调用此工具，不可编造。")
def RAGTool(question: str) -> str:
    rag_data = {
        "公司年假政策": "员工入职满1年5天年假，满3年10天，满5年15天。",
        "产品A的价格": "产品A基础版999元/年，高级版1999元/年，企业版需定制。",
        "加班补贴标准": "工作日20元/小时，周末30元/小时，法定节假日50元/小时。"
    }
    for key, value in rag_data.items():
        if key in question:
            return value
    return f"知识库中未找到「{question}」的信息，请联系相关部门。"


# 2. 整理工具列表
tools = [WeatherTool, CalculatorTool, RAGTool]

# 3. 初始化 LLM（通义千问 qwen-turbo）
llm = ChatOpenAI(
    api_key=os.getenv("DASHSCOPE_API_KEY"),
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    model="qwen-turbo",
    temperature=0  # 决策严谨，避免随机输出
)

# 4. 定义提示词模板（含多轮对话上下文）
prompt_template = """
你是帮助用户解决问题的智能Agent，可使用以下工具：
{tools}
可用工具名称：{tool_names}

{chat_history}  # 多轮对话上下文：显示历史问答，帮助理解当前问题

使用规则：
1. 先结合历史对话理解当前问题：无需工具则直接回答，需要工具则选对应工具（根据工具描述匹配）；
2. 调用工具必须用格式：```json{{"name":"工具名","parameters":{{"参数名":"参数值"}}}}```（参数需齐全，无空格）；
3. 拿到工具结果后，判断是否需要继续调用工具，或整理结果回答（若有历史上下文，需关联之前的信息）。

用户当前问题：{input}
思考过程：{agent_scratchpad}
"""

# 5. 创建提示词（包含所有必要变量）
prompt = PromptTemplate(
    template=prompt_template,
    input_variables=["tools", "tool_names", "chat_history", "input", "agent_scratchpad"]
)

print("提示词：\n" , prompt)

# 6. 创建 ZeroShotAgent（工厂方法）
agent = ZeroShotAgent.from_llm_and_tools(
    llm=llm,
    tools=tools,
    prompt=prompt,
    verbose=True
)

# 7. 初始化多轮对话记忆 + 创建 Agent 执行器
memory = ConversationBufferMemory(memory_key="chat_history")  # 记忆键名与提示词一致
agent_executor = AgentExecutor(
    agent=agent,
    tools=tools,
    verbose=True,
    handle_parsing_errors=True,
    max_iterations=5,
    memory=memory  # 绑定记忆
)

# 8. 测试场景（用 invoke 替代 run，支持多轮对话）
if __name__ == "__main__":
    # 场景1：无需工具（AI定义）
    print("=" * 50)
    print("场景1：无需工具 → 什么是人工智能？")
    result1 = agent_executor.invoke({"input": "什么是人工智能？用简单的话解释。"})
    print("答案：", result1["output"])

    # 场景2：单个工具（查北京天气）
    print("\n" + "=" * 50)
    print("场景2：单个工具 → 查询北京今天的天气")
    result2 = agent_executor.invoke({"input": "查询北京今天的天气。"})
    print("答案：", result2["output"])

    # 场景3：多个工具（上海天气+穿搭+预算，支持追问）
    print("\n" + "=" * 50)
    print("场景3：多个工具 → 上海天气+穿搭+预算（首次查询）")
    result3 = agent_executor.invoke({"input": "查上海今天的天气，推荐穿搭，再算预算（T恤150元，牛仔裤200元）。"})
    print("答案：", result3["output"])

    # 场景3 追问：加鞋子后的预算（测试多轮对话）
    print("\n" + "=" * 50)
    print("场景3 追问：加一双100元的鞋子，总预算多少？")
    result3_2 = agent_executor.invoke({"input": "如果再加一双100元的鞋子，总预算多少？"})
    print("答案：", result3_2["output"])

    # 场景4：RAG工具（年假政策）
    print("\n" + "=" * 50)
    print("场景4：RAG工具 → 公司年假政策")
    result4 = agent_executor.invoke({"input": "公司员工入职满3年有多少天年假？"})
    print("答案：", result4["output"])